// =============================================================================
// PROJECT CHRONO - http://projectchrono.org
//
// Copyright (c) 2014 projectchrono.org
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be found
// in the LICENSE file at the top level of the distribution and at
// http://projectchrono.org/license-chrono.txt.
//
// =============================================================================
// Authors: Alessandro Tasora
// =============================================================================
//
// Demo on low-level functionality for time integration of differential equations.
//
// =============================================================================

#include <cmath>

#include "chrono/core/ChDataPath.h"
#include "chrono/physics/ChSystemNSC.h"
#include "chrono/solver/ChDirectSolverLS.h"

#include "chrono_postprocess/ChGnuPlot.h"

#include "chrono_thirdparty/filesystem/path.h"

using namespace chrono;
using namespace postprocess;

void example1(const std::string& out_dir) {
    std::cout << " Example 1: integrate dx/dt=e^t " << std::endl;

    // Define a class inherited from ChIntegrable,
    // it will represent the differential equations
    // by implementing the StateSolve() function, and few other interfaces:
    class MyIntegrable : public ChIntegrable {
      private:
      public:
        MyIntegrable() {}

        // the number of coordinates in the state:
        virtual unsigned int GetNumCoordsPosLevel() override { return 0; }
        virtual unsigned int GetNumCoordsVelLevel() override { return 1; }
        virtual unsigned int GetNumCoordsAccLevel() override { return 0; }

        // compute  dy/dt=f(y,t)
        virtual bool StateSolve(ChStateDelta& dydt,                // result: computed dy/dt
                                ChVectorDynamic<>& L,              // result: computed Lagrange multipliers
                                const ChState& y,                  // current state y
                                const double T,                    // current time T
                                const double dt,                   // timestep (if needed)
                                bool force_state_scatter,          // if false, y and T are not scattered to the system
                                bool full_update,                  // if true, perform a full update during scatter
                                ChLumpingParms* lumping = nullptr  // if not null, uses lumped masses to avoid
                                                                   // inverting a mass matrix. Not significant here.
                                ) override {
            if (force_state_scatter)
                StateScatter(y, T, full_update);  // state -> system   (not needed here, btw.)

            dydt(0) = std::exp(T);  // dx/dt=e^t

            return true;
        }
    };

    // File to dump results
    std::string logfile = out_dir + "/log_timestepper_1.dat";
    std::ofstream log_file1(out_dir + "/log_timestepper_1.dat");

    // Create and object from your custom integrable class:
    MyIntegrable mintegrable;

    // Create a time-integrator, class: Euler explicit
    ChTimestepperEulerExplicitIorder mystepper(&mintegrable);

    // Execute the time integration
    while (mystepper.GetTime() < 4) {
        mystepper.Advance(0.1);

        double exact_solution = std::exp(mystepper.GetTime()) - 1;
        std::cout << " T = " << mystepper.GetTime() << "  x=" << mystepper.GetState()(0)
                  << "  x_exact=" << exact_solution << std::endl;
        log_file1 << mystepper.GetTime() << ", " << mystepper.GetState()(0) << ", " << exact_solution << std::endl;
    }

    // Plot results
    std::string gplfile = out_dir + "/tmp_timestepping_1.gpl";
    ChGnuPlot mplot(gplfile);
    mplot.SetGrid();
    mplot.SetTitle("Integrate dx/dt=e^t ");
    mplot.SetLabelX("t");
    mplot.SetLabelY("x");
    mplot.Plot(logfile, 1, 2, "Euler explicit", " with lines lt -1 lw 2");
    mplot.Plot(logfile, 1, 3, "Exact, analytical", " with lines lt 2 lw 2");
}

void example2(const std::string& out_dir) {
    std::cout
        << std::endl
        << std::endl
        << "Example 2: integrate 2nd order oscillator: M*ddx/dtdt + R*dx/dt + K*w = F(t) with a 1st order integrator. "
        << std::endl;

    // Define a class inherited from ChIntegrable,
    // it will represent the differential equations
    // by implementing the StateSolve() function, and few other interfaces:
    class MyIntegrable : public ChIntegrable {
      private:
        double m_M;
        double m_K;
        double m_R;
        double m_T;
        double m_x;
        double m_v;

      public:
        MyIntegrable() {
            m_T = 0;
            m_M = 10;
            m_K = 30;
            m_R = 1;
            m_x = 0;
            m_v = 0;
        }

        // the number of coordinates in the state:
        virtual unsigned int GetNumCoordsPosLevel() override { return 1; }
        virtual unsigned int GetNumCoordsVelLevel() override { return 1; }
        virtual unsigned int GetNumCoordsAccLevel() override { return 1; }

        // system -> state
        virtual void StateGather(ChState& y, double& T) override {
            y(0) = m_x;
            y(1) = m_v;
            T = m_T;
        }

        // state -> system
        virtual void StateScatter(const ChState& y, const double T, bool full_update) override {
            m_x = y(0);
            m_v = y(1);
            m_T = T;
        }

        // compute  dy/dt=f(y,t)
        virtual bool StateSolve(ChStateDelta& dydt,                // result: computed dy/dt
                                ChVectorDynamic<>& L,              // result: computed Lagrange multipliers
                                const ChState& y,                  // current state y
                                const double t,                    // current time T
                                const double dt,                   // timestep (if needed)
                                bool force_state_scatter,          // if false, y and T are not scattered to the system
                                bool full_update,                  // if true, perform a full update during scatter
                                ChLumpingParms* lumping = nullptr  // if not null, uses lumped masses to avoid
                                                                   // inverting a mass matrix. Not significant here.
                                ) override {
            if (force_state_scatter)
                StateScatter(y, t, full_update);

            double F = std::cos(t * 20) * 2;

            dydt(0) = m_v;                                       // speed
            dydt(1) = (1. / m_M) * (F - m_K * m_x - m_R * m_v);  // acceleration

            return true;
        }
    };

    // File to dump results
    std::string logfile = out_dir + "/log_timestepper_2.dat";
    std::ofstream log_file2(out_dir + "/log_timestepper_2.dat");

    // Try integrator Euler explicit

    // Create and object from your custom integrable class:
    MyIntegrable mintegrable;
    // Create a time-integrator:
    ChTimestepperEulerExplicitIorder mystepper(&mintegrable);

    // Try integrator Runge Kutta 4st  explicit

    // Create and object from your custom integrable class:
    MyIntegrable mintegrable_rk;
    // Create a time-integrator, class: Runge Kutta 4 explicit
    ChTimestepperRungeKutta mystepper_rk(&mintegrable_rk);

    // Execute the time integration
    while (mystepper.GetTime() < 1) {
        mystepper.Advance(0.01);
        mystepper_rk.Advance(0.01);

        std::cout << " T = " << mystepper.GetTime() << "  x=" << mystepper.GetState()(0)
                  << "  v=" << mystepper.GetState()(1) << std::endl;
        log_file2 << mystepper.GetTime() << ", " << mystepper.GetState()(0) << ", " << mystepper.GetState()(1) << ", "
                  << mystepper_rk.GetState()(0) << ", " << mystepper_rk.GetState()(1) << std::endl;
    }

    // Plot results
    std::string gplfile = out_dir + "/tmp_timestepping_2.gpl";
    ChGnuPlot mplot(gplfile);
    mplot.SetGrid();
    mplot.SetTitle("Integrate 2nd order oscillator with 1st order timestepper");
    mplot.SetLabelX("t");
    mplot.SetLabelY("x, v");
    mplot.Plot(logfile, 1, 2, "Euler exp. x", " with lines");
    mplot.Plot(logfile, 1, 3, "Euler exp. v", " with lines");
    mplot.Plot(logfile, 1, 4, "RungeKutta x", " with lines");
    mplot.Plot(logfile, 1, 5, "RungeKutta v", " with lines");
}

void example3(const std::string& out_dir) {
    std::cout << std::endl
              << std::endl
              << "Example 3: integrate 2nd order oscillator: M*ddx/dtdt + R*dx/dt + K*w = F(t) with a 2nd "
                 "order integrator. "
              << std::endl;

    // Define a class inherited from ChIntegrableIIorder,
    // it will represent the differential equations
    // by implementing the StateSolveA() function, and few other interfaces.
    // Compared to previous example 2, this is inherited from ChIntegrableIIorder,
    // so one can use more advanced integrators such as ChTimestepperEulerSemiImplicit,
    // that can exploit the II order nature of the problem. Anyway, also all I order
    // integrators can still be used.
    class MyIntegrable : public ChIntegrableIIorder {
      private:
        double m_M;
        double m_K;
        double m_R;
        double m_T;
        double m_x;
        double m_v;

      public:
        MyIntegrable() {
            m_T = 0;
            m_M = 10;
            m_K = 30;
            m_R = 1;
            m_x = 0;
            m_v = 0;
        }

        // the number of coordinates in the state, x position part:
        virtual unsigned int GetNumCoordsPosLevel() override { return 1; }
        virtual unsigned int GetNumCoordsVelLevel() override { return 1; }
        virtual unsigned int GetNumCoordsAccLevel() override { return 1; }

        // system -> state
        virtual void StateGather(ChState& x, ChStateDelta& v, double& T) override {
            x(0) = m_x;
            v(0) = m_v;
            T = m_T;
        };

        // state -> system
        virtual void StateScatter(const ChState& x, const ChStateDelta& v, const double T, bool full_update) override {
            m_x = x(0);
            m_v = v(0);
            m_T = T;
        };

        // compute  dy/dt=f(y,t)
        virtual bool StateSolveA(ChStateDelta& dvdt,                // result: computed accel. a = dv/dt
                                 ChVectorDynamic<>& L,              // result: computed Lagrange multipliers
                                 const ChState& x,                  // current state, x
                                 const ChStateDelta& v,             // current state, v
                                 const double T,                    // current time T
                                 const double dt,                   // timestep (if needed)
                                 bool force_state_scatter,          // if false, y and T are not scattered to the system
                                 bool full_update,                  // if true, perform a full update during scatter
                                 ChLumpingParms* lumping = nullptr  // if not null, uses lumped masses to avoid
                                                                    // inverting a mass matrix. Not significant here.
                                 ) override {
            if (force_state_scatter)
                StateScatter(x, v, T, full_update);

            double F = std::cos(T * 5) * 2;
            dvdt(0) = (1. / m_M) * (F - m_K * m_x - m_R * m_v);

            return true;
        }
    };

    // Create a file to dump results
    std::string logfile = out_dir + "/log_timestepper_3.dat";
    std::string logfilename = out_dir + "/log_timestepper_3.dat";
    std::ofstream log_file3(logfilename);

    // Create and object from your custom integrable class:
    MyIntegrable mintegrable1;
    MyIntegrable mintegrable2;
    MyIntegrable mintegrable3;

    // Create few time-integrators to be compared:
    ChTimestepperRungeKutta mystepper1(&mintegrable1);
    ChTimestepperEulerExplicitIIorder mystepper2(&mintegrable2);
    ChTimestepperEulerSemiImplicit mystepper3(&mintegrable3);

    // Execute the time integration
    while (mystepper1.GetTime() < 4) {
        mystepper1.Advance(0.01);
        mystepper2.Advance(0.01);
        mystepper3.Advance(0.01);

        std::cout << "T = " << mystepper1.GetTime() << "  x=" << mystepper1.GetState()(0)
                  << "  v=" << mystepper1.GetState()(1) << std::endl;
        log_file3 << mystepper1.GetTime() << ", " << mystepper1.GetState()(0) << ", " << mystepper1.GetState()(1)
                  << ", " << mystepper2.GetStatePos()(0) << ", " << mystepper2.GetStateVel()(0) << ", "
                  << mystepper3.GetStatePos()(0) << ", " << mystepper3.GetStateVel()(0) << std::endl;
    }

    // Plot results
    std::string gplfile = out_dir + "/tmp_timestepping_3.gpl";
    ChGnuPlot mplot(gplfile);
    mplot.SetGrid();
    mplot.SetTitle("Integrate 2nd order oscillator with 2nd order timestepper");
    mplot.SetLabelX("t");
    mplot.SetLabelY("x");
    mplot.Plot(logfilename, 1, 2, "RungeKutta", " with lines");
    mplot.Plot(logfilename, 1, 4, "Euler exp. IIorder", " with lines");
    mplot.Plot(logfilename, 1, 6, "Euler semi-implicit", " with lines");
}

void example4(const std::string& out_dir) {
    std::cout << std::endl
              << std::endl
              << "Example 4: integrate 2nd order oscillator: M*ddx/dtdt + R*dx/dt + K*w = F(t) with an "
                 "implicit integrator "
              << std::endl;

    // Define a class inherited from ChIntegrableIIorder,
    // it will represent the differential equations
    // by implementing the interfaces to implicit solvers.
    //  We assume   M*a = F(x,v,t)

    class MyIntegrable : public ChIntegrableIIorder {
      private:
        double m_M;
        double m_K;
        double m_R;
        double m_T;
        double m_x;
        double m_v;
        double m_a;

      public:
        MyIntegrable() {
            m_T = 0;
            m_M = 1;
            m_K = 30;
            m_R = 0;
            m_x = 0;
            m_v = 0.6;
            m_a = 0;
        }

        // the number of coordinates in the state, x position part:
        virtual unsigned int GetNumCoordsPosLevel() override { return 1; }
        virtual unsigned int GetNumCoordsVelLevel() override { return 1; }
        virtual unsigned int GetNumCoordsAccLevel() override { return 1; }

        // system -> state
        virtual void StateGather(ChState& x, ChStateDelta& v, double& T) override {
            x(0) = m_x;
            v(0) = m_v;
            T = m_T;
        };

        // state -> system
        virtual void StateScatter(const ChState& x, const ChStateDelta& v, const double T, bool full_update) override {
            m_x = x(0);
            m_v = v(0);
            m_T = T;
        };

        // gather/scatter of accelerations not needed for some solvers (exEuler linearized) but needed for others, ex.
        // HHT
        virtual void StateGatherAcceleration(ChStateDelta& a) override { a(0) = m_a; }
        virtual void StateScatterAcceleration(const ChStateDelta& a) override { m_a = a(0); }

        // compute  dy/dt=f(y,t)
        // (this function is optional: if not implemented the integrator can solve
        // for acceleration also using StateSolveCorrection, although a bit less efficient)
        virtual bool StateSolveA(
            ChStateDelta& dvdt,                // result: computed accel. a=dv/dt
            ChVectorDynamic<>& L,              // result: computed Lagrange multipliers
            const ChState& x,                  // current state, x
            const ChStateDelta& v,             // current state, v
            const double T,                    // current time T
            const double dt,                   // timestep (if needed)
            bool force_state_scatter,          // if false, y and T are not scattered to the system
            bool full_update,                  // if true, perform a full update during scatter
            ChLumpingParms* lumping = nullptr  // if not null, uses lumped masses to avoid inverting a mass matrix,
                                               // and uses penalty for constraints. Not significant here.
            ) override {
            if (force_state_scatter)
                StateScatter(x, v, T, full_update);
            double F = std::sin(m_T * 20) * 0.02;
            dvdt(0) = (1. / m_M) * (F - m_K * m_x - m_R * m_v);

            return true;
        }

        // Compute the correction with linear system
        //  Dv = [ c_a*M + c_v*dF/dv + c_x*dF/dx ]^-1 * R
        virtual bool StateSolveCorrection(ChStateDelta& Dv,             // result: computed Dv
                                          ChVectorDynamic<>& L,         // result: computed Lagrange multipliers
                                          const ChVectorDynamic<>& R,   // the R residual
                                          const ChVectorDynamic<>& Qc,  // the Qc residual
                                          const double c_a,             // the factor in c_a*M
                                          const double c_v,             // the factor in c_v*dF/dv
                                          const double c_x,             // the factor in c_x*dF/dv
                                          const ChState& x,             // current state, x part
                                          const ChStateDelta& v,        // current state, v part
                                          const double T,               // current time T
                                          bool force_state_scatter,     // if true, scatter x and v to the system
                                          bool full_update,             // if true, perform a full update during scatter
                                          bool call_setup,              // if true, call the solver's Setup function
                                          bool call_analyze  // if true, call the solver's Setup analyze phase
                                          ) override {
            if (force_state_scatter)
                this->StateScatter(x, v, T, full_update);

            Dv(0) = R(0) * 1.0 / (c_a * m_M + c_v * (-m_R) + c_x * (-m_K));

            return true;
        }

        //    R += c*F
        void LoadResidual_F(ChVectorDynamic<>& R,  // result: the R residual, R += c*F
                            const double c         // a scaling factor
                            ) override {
            R(0) += c * (std::sin(m_T * 20) * 0.02 - m_K * m_x - m_R * m_v);
        };

        //    R += c*M*w
        void LoadResidual_Mv(ChVectorDynamic<>& R,        // result: the R residual, R += c*M*v
                             const ChVectorDynamic<>& w,  // the w vector
                             const double c               // a scaling factor
                             ) override {
            R(0) += c * m_M * w(0);
        };

        // nothing to do here- no constraints
        virtual void LoadResidual_CqL(ChVectorDynamic<>& R,        // result: the R residual, R += c*Cq'*L
                                      const ChVectorDynamic<>& L,  // the L vector
                                      const double c               // a scaling factor
                                      ) override {}

        // nothing to do here- no constraints
        virtual void LoadConstraint_C(ChVectorDynamic<>& Qc,        // result: the Qc residual, Qc += c*C
                                      const double c,               // a scaling factor
                                      const bool do_clamp = false,  // enable optional clamping of Qc
                                      const double mclam = 1e30     // clamping value
                                      ) override {}

        // nothing to do here- no constraints
        virtual void LoadConstraint_Ct(ChVectorDynamic<>& Qc,  // result: the Qc residual, Qc += c*Ct
                                       const double c          // a scaling factor
                                       ) override {}
    };

    // Create a file to dump results
    std::string logfilename = out_dir + "/log_timestepper_4.dat";
    std::ofstream log_file4(logfilename);

    // Create and object from your custom integrable class:
    MyIntegrable mintegrable1;
    MyIntegrable mintegrable2;
    MyIntegrable mintegrable3;
    MyIntegrable mintegrable4;
    MyIntegrable mintegrable5;
    MyIntegrable mintegrable6;
    MyIntegrable mintegrable7;
    MyIntegrable mintegrable8;

    // Create few time-integrators to be compared:
    ChTimestepperEulerImplicit mystepper1(&mintegrable1);
    ChTimestepperTrapezoidal mystepper2(&mintegrable2);
    ChTimestepperEulerExplicitIIorder mystepper3(&mintegrable3);
    ChTimestepperHHT mystepper4(&mintegrable4);
    mystepper4.SetAlpha(0);  // HHT with no dissipation -> trapezoidal
    ChTimestepperHHT mystepper5(&mintegrable5);
    mystepper5.SetAlpha(-0.33);  // HHT with max dissipation
    ChTimestepperNewmark mystepper6(&mintegrable6);
    mystepper6.SetGammaBeta(0.5, 0.25);  // Newmark as const accel. method
    ChTimestepperNewmark mystepper7(&mintegrable7);
    mystepper7.SetGammaBeta(0.5, 1 / 6);  // Newmark as linear accel. method
    ChTimestepperNewmark mystepper8(&mintegrable8);
    mystepper8.SetGammaBeta(1.0, 0.25);  // Newmark with max numerical damping

    // Execute the time integration
    while (mystepper1.GetTime() < 4) {
        double timestep = 0.05;
        mystepper1.Advance(timestep);
        mystepper2.Advance(timestep);
        mystepper3.Advance(timestep);
        mystepper4.Advance(timestep);
        mystepper5.Advance(timestep);
        mystepper6.Advance(timestep);
        mystepper7.Advance(timestep);
        mystepper8.Advance(timestep);

        std::cout << "T = " << mystepper1.GetTime() << "  x=" << mystepper1.GetStatePos()(0)
                  << "  v=" << mystepper1.GetStateVel()(0) << std::endl;
        log_file4 << mystepper1.GetTime() << ", " << mystepper1.GetStatePos()(0) << ", " << mystepper1.GetStateVel()(0)
                  << ", " << mystepper2.GetStatePos()(0) << ", " << mystepper2.GetStateVel()(0) << ", "
                  << mystepper3.GetStatePos()(0) << ", " << mystepper3.GetStateVel()(0) << ", "
                  << mystepper4.GetStatePos()(0) << ", " << mystepper4.GetStateVel()(0) << ", "
                  << mystepper5.GetStatePos()(0) << ", " << mystepper5.GetStateVel()(0) << ", "
                  << mystepper6.GetStatePos()(0) << ", " << mystepper6.GetStateVel()(0) << ", "
                  << mystepper7.GetStatePos()(0) << ", " << mystepper7.GetStateVel()(0) << ", "
                  << mystepper8.GetStatePos()(0) << ", " << mystepper8.GetStateVel()(0) << std::endl;
    }

    // Plot results
    std::string gplfile = out_dir + "/tmp_timestepping_4.gpl";
    ChGnuPlot mplot(gplfile);
    mplot.SetGrid();
    mplot.SetTitle("Test: oscillator with implicit integrators");
    mplot.SetLabelX("t");
    mplot.SetLabelY("x");
    mplot.Plot(logfilename, 1, 2, "Euler implicit", " with lines");
    mplot.Plot(logfilename, 1, 4, "Trapezoidal", " with lines");
    mplot.Plot(logfilename, 1, 6, "Euler expl.IIorder", " with lines");
    mplot.Plot(logfilename, 1, 8, "HHT alpha=0", " with lines dt 2");
    mplot.Plot(logfilename, 1, 10, "HHT alpha=-0.33", " with lines dt 2");
    mplot.Plot(logfilename, 1, 12, "Newmark g=0.5, b=1/4", " with lines dt 4");
    mplot.Plot(logfilename, 1, 14, "Newmark g=0.5, b=1/6", " with lines dt 4");
    mplot.Plot(logfilename, 1, 16, "Newmark g=1.0, b=1/4", " with lines dt 4");
}

void example5(const std::string& out_dir) {
    std::cout << "\n\nExample 5: integrate pendulum DAE " << std::endl;

    // A) - a minimalistic pendulum DAE:
    //
    // Define a class inherited from ChIntegrableIIorder,
    // it will represent the differential equations
    // by implementing the interfaces to implicit solvers.
    //  We assume   M*a = F(x,v,t)
    //            C(x,t)=0;

    class MyIntegrable : public ChIntegrableIIorder {
      private:
        double m_M;
        double m_K;
        double m_R;
        double m_T;
        double m_px;
        double m_py;
        double m_vx;
        double m_vy;
        double m_ax;
        double m_ay;
        double m_length;
        double m_reaction;

      public:
        MyIntegrable() {
            m_length = 5;
            m_M = 2;
            m_K = 0;
            m_R = 0;
            m_T = 0;
            m_px = 0;
            m_py = -m_length;
            m_vx = 0.8;
            m_vy = 0;
            m_ax = 0;
            m_ay = 0;
            m_reaction = 0;
        }

        // Number of coordinates in the state, position level
        virtual unsigned int GetNumCoordsPosLevel() override { return 2; }
        // Number of coordinates in the state, velocity level
        virtual unsigned int GetNumCoordsVelLevel() override { return 2; }
        // Number of coordinates in the state, acceleration level
        virtual unsigned int GetNumCoordsAccLevel() override { return 2; }

        // Return the number of Lagrange multipliers (constraints)
        virtual unsigned int GetNumConstraints() override { return 1; }

        // system -> state
        virtual void StateGather(ChState& x, ChStateDelta& v, double& T) override {
            x(0) = m_px;
            x(1) = m_py;
            v(0) = m_vx;
            v(1) = m_vy;
            T = m_T;
        };

        // state -> system
        virtual void StateScatter(const ChState& x, const ChStateDelta& v, const double T, bool full_update) override {
            m_px = x(0);
            m_py = x(1);
            m_vx = v(0);
            m_vy = v(1);
            m_T = T;
        };

        virtual void StateGatherAcceleration(ChStateDelta& a) override {
            a(0) = m_ax;
            a(1) = m_ay;
        }
        virtual void StateScatterAcceleration(const ChStateDelta& a) override {
            m_ax = a(0);
            m_ay = a(1);
        }

        // Some timesteppers exploit persistence of reaction information
        virtual void StateGatherReactions(ChVectorDynamic<>& L) override { L(0) = m_reaction; };
        virtual void StateScatterReactions(const ChVectorDynamic<>& L) override { m_reaction = L(0); };

        // Compute the correction with linear system
        //  | Dv| = [ c_a*M + c_v*dF/dv + c_x*dF/dx    Cq']^-1 * | R |
        //  |-Dl|   [   Cq                              0 ]      |-Qc|
        virtual bool StateSolveCorrection(ChStateDelta& Dv,             // result: computed Dv
                                          ChVectorDynamic<>& Dl,        // result: computed Dl Lagrange multipliers
                                          const ChVectorDynamic<>& R,   // the R residual
                                          const ChVectorDynamic<>& Qc,  // the Qc residual
                                          const double c_a,             // the factor in c_a*M
                                          const double c_v,             // the factor in c_v*dF/dv
                                          const double c_x,             // the factor in c_x*dF/dv
                                          const ChState& x,             // current state, x part
                                          const ChStateDelta& v,        // current state, v part
                                          const double T,               // current time T
                                          bool force_state_scatter,     // if true, scatter x and v to the system
                                          bool full_update,             // if true, perform a full update during scatter
                                          bool call_setup,              // if true, call the solver's Setup function
                                          bool call_analyze  // if true, call the solver's Setup analyze phase
                                          ) override {
            if (force_state_scatter)
                this->StateScatter(x, v, T, full_update);

            ChVector3d dirpend(-m_px, -m_py, 0);
            dirpend.Normalize();
            ChVectorDynamic<> b(3);
            b(0) = R(0);
            b(1) = R(1);
            b(2) = -Qc(0);  // note assume input Qc has no minus sign, so flip sign here
            ChMatrixDynamic<> A(3, 3);
            A.setZero();
            A(0, 0) = c_a * m_M + c_v * (-m_R) + c_x * (-m_K);
            A(1, 1) = c_a * m_M;
            A(0, 2) = dirpend.x();
            A(1, 2) = dirpend.y();
            A(2, 0) = dirpend.x();
            A(2, 1) = dirpend.y();
            ChVectorDynamic<> w = A.colPivHouseholderQr().solve(b);
            Dv(0) = w(0);
            Dv(1) = w(1);
            Dl(0) = -w(2);  // note assume result sign in multiplier is flipped, to return Dl and not w=-Dl

            return true;
        }

        // Adds the lumped mass to a Md vector. This method is OPTIONAL, and needed only
        // if you want to use an explicit integrator with SetDiagonalLumpingON.
        virtual void LoadLumpedMass_Md(ChVectorDynamic<>& Md,  // diagonal of the lumped mass matrix
                                       double& err,            // not touched if lumping does not introduce errors
                                       const double c          // a scaling factor
                                       ) override {
            Md(0) = m_M;
            Md(1) = m_M;
        }

        //    R += c*F
        void LoadResidual_F(ChVectorDynamic<>& R,  // result: the R residual, R += c*F
                            const double c         // a scaling factor
                            ) override {
            R(0) += c * (-m_K * m_px - m_R * m_vx);
            R(1) += c * -9.8 * m_M;  // vertical force
        };

        //    R += c*M*w
        void LoadResidual_Mv(ChVectorDynamic<>& R,        // result: the R residual, R += c*M*v
                             const ChVectorDynamic<>& w,  // the w vector
                             const double c               // a scaling factor
                             ) override {
            R(0) += c * m_M * w(0);
            R(1) += c * m_M * w(1);
        };

        //   R += Cq'*l
        virtual void LoadResidual_CqL(ChVectorDynamic<>& R,        // result: the R residual, R += c*Cq'*L
                                      const ChVectorDynamic<>& L,  // the L vector
                                      const double c               // a scaling factor
                                      ) override {
            ChVector3d dirpend(-m_px, -m_py, 0);
            dirpend.Normalize();
            R(0) += c * dirpend.x() * L(0);
            R(1) += c * dirpend.y() * L(0);
        };

        //  Qc += c * C
        virtual void LoadConstraint_C(ChVectorDynamic<>& Qc,        // result: the Qc residual, Qc += c*C
                                      const double c,               // a scaling factor
                                      const bool do_clamp = false,  // enable optional clamping of Qc
                                      const double mclam = 1e30     // clamping value
                                      ) override {
            ChVector3d distpend(-m_px, -m_py, 0);
            Qc(0) += c * (-distpend.Length() + m_length);
        };

        // nothing to do here- no rheonomic part
        virtual void LoadConstraint_Ct(ChVectorDynamic<>& Qc,  // result: the Qc residual, Qc += c*Ct
                                       const double c          // a scaling factor
                                       ) override {}
    };

    // Create a file to dump results
    std::string logfilename5 = out_dir + "/log_timestepper_5.dat";
    std::string logfilename5r = out_dir + "/log_timestepper_5r.dat";
    std::string logfilename5e = out_dir + "/log_timestepper_5e.dat";
    std::string logfilename5er = out_dir + "/log_timestepper_5er.dat";
    std::ofstream log_file5(logfilename5);
    std::ofstream log_file5r(logfilename5r);
    std::ofstream log_file5e(logfilename5e);
    std::ofstream log_file5er(logfilename5er);

    // Create and object from your custom integrable class:
    MyIntegrable mintegrable1;
    MyIntegrable mintegrable2;
    MyIntegrable mintegrable3;
    MyIntegrable mintegrable4;
    MyIntegrable mintegrable5;
    MyIntegrable mintegrable6;
    MyIntegrable mintegrable7;

    // Create few time-integrators to be compared:
    ChTimestepperEulerImplicitLinearized mystepper1(&mintegrable1);
    ChTimestepperEulerImplicit mystepper2(&mintegrable2);
    ChTimestepperTrapezoidal mystepper3(&mintegrable3);
    ChTimestepperHHT mystepper4(&mintegrable4);
    mystepper4.SetAlpha(0);  // HHT with no dissipation -> trapezoidal
    mystepper4.SetStepControl(false);
    // mystepper4.SetVerbose(true);
    ChTimestepperHHT mystepper5(&mintegrable5);
    mystepper5.SetAlpha(-0.3);  // HHT with dissipation
    mystepper5.SetStepControl(false);
    // mystepper5.SetVerbose(true);
    ChTimestepperNewmark mystepper6(&mintegrable6);
    mystepper6.SetGammaBeta(
        0.5, 0.25);  // Newmark, Gamma: in [1/2, 1] where 1/2 no damping, beta in [0,1]. For (0.5, 0.25) -> trapezoidal
    // mystepper6.SetVerbose(true);

    // ChTimestepperEulerExplicitIIorder mystepper7(&mintegrable7);
    ChTimestepperRungeKutta mystepper7(&mintegrable7);
    mystepper7.SetDiagonalLumpingON(20000);  // this avoids calling the linear solver completely, even with constraints.

    // B) - same pendulum, but multibody:
    //
    // Ok, let's create the same pendulum but using the multibody toolset of Chrono,
    // that is using the ChBody, ChLinkLockRevolute and ChSystem classes:

    ChSystemNSC sys;
    auto my_body_A = chrono_types::make_shared<ChBody>();
    auto my_body_B = chrono_types::make_shared<ChBody>();
    sys.AddBody(my_body_A);
    sys.AddBody(my_body_B);

    my_body_A->SetFixed(true);
    my_body_B->SetMass(2.0);
    my_body_B->SetInertiaXX(ChVector3d(1e-7, 1e-7, 1e-7));  // to approximate point-like mass as in MyIntegrable
    my_body_B->SetPos(ChVector3d(0, -5, 0));
    my_body_B->SetPosDt(ChVector3d(0.8, 0, 0));

    auto my_link_AB = chrono_types::make_shared<ChLinkLockRevolute>();
    my_link_AB->Initialize(my_body_A, my_body_B, ChFrame<>());
    sys.AddLink(my_link_AB);

    // use a precise linear solver
    auto msolver = chrono_types::make_shared<ChSolverSparseQR>();
    sys.SetSolver(msolver);
    // use the HHT timestepper to compare with HHT in previous MyIntegrable simple case
    auto mstepper4b = chrono_types::make_shared<ChTimestepperHHT>(&sys);
    mstepper4b->SetAlpha(-0.3);  // HHT dissipation
    sys.SetTimestepper(mstepper4b);

    // Execute the time integration with the implicit integrators
    while (mystepper1.GetTime() < 12) {
        double timestep = 0.05;
        mystepper1.Advance(timestep);
        mystepper2.Advance(timestep);
        mystepper3.Advance(timestep);
        mystepper4.Advance(timestep);
        mystepper5.Advance(timestep);
        mystepper6.Advance(timestep);
        sys.DoStepDynamics(timestep);

        std::cout << "T = " << mystepper1.GetTime() << "  x=" << mystepper1.GetStatePos()(0)
                  << "  y=" << mystepper1.GetStatePos()(1) << std::endl;
        std::cout << "  = " << mystepper1.GetTime() << "  x=" << my_body_B->GetPos().x()
                  << "  y=" << my_body_B->GetPos().y() << std::endl;
        log_file5 << mystepper1.GetTime() << ", " << mystepper1.GetStatePos()(0) << ", " << mystepper1.GetStatePos()(1)
                  << ", " << mystepper1.GetStateVel()(0) << ", " << mystepper1.GetStateVel()(1) << ", "
                  << mystepper2.GetStatePos()(0) << ", " << mystepper2.GetStatePos()(1) << ", "
                  << mystepper2.GetStateVel()(0) << ", " << mystepper2.GetStateVel()(1) << ", "
                  << mystepper3.GetStatePos()(0) << ", " << mystepper3.GetStatePos()(1) << ", "
                  << mystepper3.GetStateVel()(0) << ", " << mystepper3.GetStateVel()(1) << ", "
                  << mystepper4.GetStatePos()(0) << ", " << mystepper4.GetStatePos()(1) << ", "
                  << mystepper4.GetStateVel()(0) << ", " << mystepper4.GetStateVel()(1) << ", "
                  << mystepper5.GetStatePos()(0) << ", " << mystepper5.GetStatePos()(1) << ", "
                  << mystepper5.GetStateVel()(0) << ", " << mystepper5.GetStateVel()(1) << ", "
                  << mystepper6.GetStatePos()(0) << ", " << mystepper6.GetStatePos()(1) << ", "
                  << mystepper6.GetStateVel()(0) << ", " << mystepper6.GetStateVel()(1) << ", "
                  << my_body_B->GetPos().x() << ", " << my_body_B->GetPos().y() << ", " << my_body_B->GetPosDt().x()
                  << ", " << my_body_B->GetPosDt().y() << std::endl;
        log_file5r << mystepper1.GetTime() << ", " << mystepper1.GetLagrangeMultipliers()(0) << ", "
                   << mystepper2.GetLagrangeMultipliers()(0) << ", " << mystepper3.GetLagrangeMultipliers()(0) << ", "
                   << mystepper4.GetLagrangeMultipliers()(0) << ", " << mystepper5.GetLagrangeMultipliers()(0) << ", "
                   << mystepper6.GetLagrangeMultipliers()(0) << ", " << my_link_AB->GetReaction2().force.y()
                   << std::endl;
    }
    // Execute the time integration with the explicit integrator,
    // hence using smaller time step
    while (mystepper7.GetTime() < 12) {
        double timestep = 0.0005;
        mystepper7.Advance(timestep);

        log_file5e << mystepper7.GetTime()
                   << ", "
                   //<< mystepper7.GetStatePos()(0) << ", " << mystepper7.GetStatePos()(1) << ", " <<
                   // mystepper7.GetStateVel()(0) << ", "
                   //<< mystepper7.GetStateVel()(1) << ", "
                   << mystepper7.GetState()(0) << ", " << mystepper7.GetState()(1) << ", " << mystepper7.GetState()(2)
                   << ", " << mystepper7.GetState()(3) << ", " << std::endl;
        log_file5er << mystepper7.GetTime() << ", " << mystepper7.GetLagrangeMultipliers()(0) << std::endl;
    }

    std::string gplfile = out_dir + "/tmp_timestepping_5.gpl";
    ChGnuPlot mplot(gplfile);
    mplot.OutputWindow(0);
    mplot.SetGrid();
    mplot.SetTitle("Test: DAE, constrained pendulum");
    mplot.SetLabelX("t");
    mplot.SetLabelY("x");
    mplot.Plot(logfilename5, 1, 2, "Euler impl. lineariz.", " with lines");
    mplot.Plot(logfilename5, 1, 6, "Euler impl.", " with lines");
    mplot.Plot(logfilename5, 1, 10, "Trapezoidal*", " with lines");
    mplot.Plot(logfilename5, 1, 14, "HHT alpha=0", " with lines dt 2");
    mplot.Plot(logfilename5, 1, 18, "HHT alpha=-0.3", " with lines dt 2");
    mplot.Plot(logfilename5, 1, 22, "Newmark g=0.5,b=0.25", " with lines dt 4");
    mplot.Plot(logfilename5, 1, 26, "HHT alpha=-0.3 in ChSystem", " with lines dt 6");
    mplot.Plot(logfilename5e, 1, 2, "Euler explicit, penalty", " with lines");

    mplot.OutputWindow(1);
    mplot.SetGrid();
    mplot.SetTitle("Test: DAE, constrained pendulum reactions");
    mplot.SetLabelX("t [s]");
    mplot.SetLabelY("R [N]");
    mplot.Plot(logfilename5r, 1, 2, "Euler impl. lineariz.", " with lines");
    mplot.Plot(logfilename5r, 1, 3, "Euler impl.", " with lines");
    mplot.Plot(logfilename5r, 1, 4, "Trapezoidal*", " with lines");
    mplot.Plot(logfilename5r, 1, 5, "HHT alpha=0", " with lines dt 2");
    mplot.Plot(logfilename5r, 1, 6, "HHT alpha=-0.3", " with lines dt 2");
    mplot.Plot(logfilename5r, 1, 7, "Newmark g=0.5,b=0.25", " with lines dt 4");
    mplot.Plot(logfilename5r, 1, 8, "HHT alpha=-0.3 in ChSystem", " with lines dt 6");
    mplot.Plot(logfilename5er, 1, 2, "Euler explicit, penalty", " with lines dt 3 lc rgb \"pink\"");

    mplot.OutputWindow(2);
    mplot.SetGrid();
    mplot.SetTitle("Test: DAE, constrained pendulum trajectory");
    mplot.SetLabelX("x");
    mplot.SetLabelY("y");
    // mplot.SetRangeX(-0.15, 0.15);
    // mplot.SetRangeY(-1.025, -0.95);
    mplot.SetCommand("set size ratio 0.5");
    mplot.Plot(logfilename5, 2, 3, "Euler impl. lineariz.", " pt 0");
    mplot.Plot(logfilename5, 6, 7, "Euler impl.", " pt 1");
    mplot.Plot(logfilename5, 10, 11, "Trapezoidal*", " pt 2");
    mplot.Plot(logfilename5, 14, 15, "HHT alpha=0", " pt 3");
    mplot.Plot(logfilename5, 18, 19, "HHT alpha=-0.2", " pt 4");
    mplot.Plot(logfilename5e, 2, 3, "Euler explicit, penalty", " with lines");
}

int main(int argc, char* argv[]) {
    std::cout << "Copyright (c) 2021 projectchrono.org\n"
              << "Chrono version: " << CHRONO_VERSION << std::endl;

    std::cout << "CHRONO demo about low-level time integration of differential equations:" << std::endl;

    // Create (if needed) output directory
    const std::string out_dir = GetChronoOutputPath() + "DEMO_TIMESTEPPER";
    if (!filesystem::create_directory(filesystem::path(out_dir))) {
        std::cout << "Error creating directory " << out_dir << std::endl;
        return 1;
    }

    example1(out_dir);
    example2(out_dir);
    example3(out_dir);
    example4(out_dir);
    example5(out_dir);
}
